import argparse

class DefaultConfig(object):
    def __init__(self,method, dataset, centralized=False, dirichlet=False, **kwargs):
        self.parser = argparse.ArgumentParser()

        # Method parameters
        if method == 'FedAvg':
            self.parser.add_argument('--method', type=str, default=method, help='Method name')
            self.parser.add_argument('--max_epoch', type=int, default=50, help='Maximum number of epochs')
            self.parser.add_argument('--alpha', type=float, default=0.3, help='Alpha')  # 第一层标签的权重
            self.parser.add_argument('--beta', type=float, default=0.7, help='Beta')  # 第二层标签的权重
            self.parser.add_argument('--gamma', type=float, default=1, help='Gamma')  # 量化损失的权重。
            self.parser.add_argument('--eta', type=int, default=0.5, help='Eta')  # 本地学习中01损失的权重。
            self.parser.add_argument('--bit', type=int, default=16, help='Final binary code length')  # 最终哈希码的长度
            self.parser.add_argument('--lr', type=float, default=0.0001, help='Learning rate')  # 模型的学习率。
            self.parser.add_argument('--train_ep', type=int, default=10, help='Number of local episodes')  # 客户端本地训练的轮次

        # Data parameters
        if dataset == 'Digits':
            self.parser.add_argument('--dataset', type=str, default='fl_digits', help='Dataset name')
            self.parser.add_argument('--data_path', type=str, default='./datasets/Digits/', help='Data path')
            self.parser.add_argument('--local_batch_size', type=int, default=128, help='Batch size')


        self.parser.add_argument('--centralized', type=bool, default=centralized, help='Centralized')  # 是否为集中式学习
        self.parser.add_argument('--Dirichlet', type=bool, default=dirichlet, help='Dirichlet')  # 是否使用Dirichlet分布

        if centralized:
            self.parser.add_argument('--num_users', type=int, default=1, help='Number of users')
        else:
            self.parser.add_argument('--num_users', type=int, default=10, help='Number of users')  # 客户端数量

        self.parser.add_argument('--device', default="cuda", type=str, help="cpu, cuda, or others")
        if dirichlet:
            self.parser.add_argument('--Beta', type=float, default=0.5, help='Dirichlet Beta')
        else:
            self.parser.add_argument('--iid',  type=bool, default=False, help='IID')
            self.parser.add_argument('--unequal', type=bool, default=False, help='Unequal')

        self.parser.add_argument('--rand_dataset', type=bool, default=True, help='The random seed')
        self.parser.add_argument('--use_gpu',  type=bool, default=True, help='Use GPU')
        self.parser.add_argument('--gpu', default=0, type=int, help='index of gpu')
        self.parser.add_argument('--val_gpu', default=0, type=str, help='index of gpu')
        self.parser.add_argument('--seed', type=int, default=42, help='Random seed')
        self.parser.add_argument('--valid',  type=bool, default=True, help='Validation') #是否每个通讯都进行一次验证
        self.parser.add_argument('--result_dir', type=str, default='result', help='Result directory') #存放训练结果的目录。

        self.opt = self.parser.parse_args(namespace=argparse.Namespace(**kwargs))

